import os
import serial
import subprocess
import time
import sys
import serial.tools.list_ports
import csv
import threading
from datetime import datetime
from time import sleep
from serial.tools import list_ports
import configparser
from colorama import init, Fore

config = configparser.ConfigParser()
config.read('config.ini')

esptool_cmd = "esptool.exe"
esp_mcu = config.get('DEFAULT', 'esp_mcu')
vid_pid = config.get('DEFAULT', 'vid_pid')
buad = config.get('DEFAULT', 'buad')
firmware_path = config.get('DEFAULT', 'firmware')

def find_devices(vid_pid):
    timeout = 30
    devices = []
    while timeout > 0:
        for p in list_ports.comports():
            if p[2].upper().startswith('USB VID:PID='+ vid_pid):
                port = p[0]
                print('Found device at port: ' + port)
                devices.append(port)
            # if p[2].upper().startswith('USB VID:PID='+ vid_pid):
            #     port = p[0]
            #     os.system("MODE %s:1200,N,8,1" %port)
            #     time.sleep(1)    
        if devices:
            #数组转字符串
            # devices = ' '.join(devices)
            print(f"共找到对应VID:PID的串口：{devices}")
            return devices
        sleep(0.1)
        timeout -= 1
    print('No test boards found')
    return devices

# 具体烧录实现
def flash_firmware(port):
    global esptool_cmd, firmware_path
    if not firmware_path or not os.path.isfile(firmware_path):
        raise ValueError(f"Firmware path is not set correctly or the file does not exist: {firmware_path}")
    cmd = [esptool_cmd, "--chip", esp_mcu, "--port", port, "--baud", buad, "write_flash", "0x0", firmware_path]
    print("{} cmd命令为:{}\n".format(port, " ".join(cmd))) 
    start_time = time.time()
    mac_address = "Unknown"
    try:
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        while True:
            output_line = process.stdout.readline()
            if output_line == '' and process.poll() is not None:
                break
            if "MAC:" in output_line:
                mac_address = output_line.split("MAC:")[1].strip()
    except Exception as e:
        print(f"Firmware flashing failed with error: {e}")
        return 0, False, mac_address
    end_time = time.time()
    flash_duration = round(end_time - start_time, 1)
    if process.poll() == 0:
        print(f"Firmware flashed successfully in {flash_duration:.2f} seconds.")
        return flash_duration, True, mac_address
    else:
        print("Firmware flashing failed.")
        return flash_duration, False, mac_address

# 失败确认后再次烧录
# def flash_firmware_with_retry(port):
#     flash_duration, flash_success, mac_address = flash_firmware(port)
#     if not flash_success:
#         input("Flash failed. Please hold the BOOT button and power on device to retry.")
#         flash_duration, flash_success, mac_address = flash_firmware(port)
#     return flash_duration, flash_success, mac_address

def append_to_csv(results, filename):
    fields = ['Timestamp', 'Port', 'MAC Address', 'Overall Result', 'Flash Duration (s)', 'Flash Result', 'Failure Reason']
    if results[-1] is None:
        results = results[:-1]
        fields = fields[:-1]
    with open(filename, 'a', newline='') as file:
        writer = csv.writer(file)
        if file.tell() == 0:
            writer.writerow(fields)
        writer.writerow(results)

# 结果统计
def flash_board(port, csv_filename):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    flash_duration, flash_success, mac_address = flash_firmware(port)
    overall_result = "PASS" if flash_success else "FAIL"
    failure_reason = None if overall_result == "PASS" else "Flash-fail" if flash_success else "Flashing failed"
    result = [timestamp, port, mac_address, overall_result, flash_duration, flash_success]
    if failure_reason:
        result.append(failure_reason)
    append_to_csv(result, csv_filename)
    return port, overall_result, flash_duration, failure_reason

# 线程函数
def worker(port, start_signal, summary_results, csv_filename):
    start_signal.wait()
    print(f"Starting on port {port}...\n")
    port, result, flash_duration, failure_reason = flash_board(port, csv_filename)
    summary_results.append((port, result, flash_duration, failure_reason))

# 显示结果
def display_summary(summary_results):
    print("==================================================")
    print("Flash Summary:")
    if not summary_results:
        print("No result were conducted.")
    else:
        for port, result, flash_duration, failure_reason in summary_results:
            color = Fore.GREEN if result == "PASS" else Fore.RED  # 根据结果选择颜色
            result_str = f"Port: {port}, Overall Result: {result}, Flash Duration: {flash_duration}s"
            if result == "FAIL":
                result_str += f", Failure Reason: {failure_reason}"
            print(f"{color}{result_str}")
    print("==================================================")
    print("End of summary.")

def main():
    init(autoreset=True)  # 初始化colorama，设置自动重置颜色
    print("\n==================================================")
    print("请插上多个待烧录产品   ")
    #print("\n")
    input("准备好后再按回车 ")
    print("==================================================")

    date_str = datetime.now().strftime("%Y-%m-%d")
    csv_filename = f"esp-flash-results-{date_str}.csv"

    while True:
        print("查找串口...    ")
        ports = find_devices(vid_pid)
        if not ports:
            print(Fore.RED +"No find board !!!")
            print(Fore.RED +"请检查VIDPID是否正确以及串口是否正常")
            input("回车后将退出程序")
            sys.exit(1)
        print(f"Ready to start flash on {ports}...")
        print(f"开始在 {ports}上进行烧录...")
        start_signal = threading.Event()
        threads = []
        summary_results = []

        for port in ports:
            t = threading.Thread(target=worker, args=(port, start_signal, summary_results, csv_filename))
            t.start()
            threads.append(t)

        start_signal.set()

        for t in threads:
            t.join()

        display_summary(summary_results)
        print("Press Enter to start a new batch.")
        input("产品插好后再按回车开始 ")
        print("==================================================\n")

if __name__ == '__main__':
    main()